import numpy as np
import cv2 as cv
import matplotlib.pyplot as plt

spr = 2
spc = 5
spn = 0
plt.figure(figsize=(13, 6))


def my_show_img(img, title, trans=None, **kwargs):
    global spn
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.axis('off')
    plt.title(title)
    if trans is not None:
        img = trans(img)
    plt.imshow(img, **kwargs)


# load
img_path = '../../../../../large_data/pic/watershed/water_coins.jpg'
img = cv.imread(img_path, cv.IMREAD_COLOR)
my_show_img(img, 'ori', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

# gray
gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
my_show_img(gray, 'gray', cmap='gray')

# bin ostu
ret, bin = cv.threshold(gray, 0, 255, cv.THRESH_BINARY_INV + cv.THRESH_OTSU)
my_show_img(bin, 'bin', cmap='gray')

# opening
kernel = np.ones((3, 3), dtype=np.uint8)
opening = cv.morphologyEx(bin, cv.MORPH_OPEN, kernel, iterations=2)
my_show_img(opening, 'opening', cmap='gray')

# sure bg
sure_bg = cv.dilate(opening, kernel, iterations=3)
my_show_img(sure_bg, 'open => sure_bg', cmap='gray')

# sure fg
dist_trans = cv.distanceTransform(opening, cv.DIST_L2, 5)
my_show_img(dist_trans, 'open => dist_trans', cmap='gray')
ret, sure_fg = cv.threshold(dist_trans, 0.7 * dist_trans.max(), 255, cv.THRESH_BINARY)
my_show_img(sure_fg, 'dist_trans => sure_fg', cmap='gray')

# unknown
sure_fg = sure_fg.astype(np.uint8)
unknown = cv.subtract(sure_bg, sure_fg)
my_show_img(unknown, 'unknown', cmap='gray')

# markers
ret, markers = cv.connectedComponents(sure_fg)
markers += 1
markers[unknown == 255] = 0
my_show_img(markers, 'markers')

# watershed
res = cv.watershed(img, markers)
print('res', res.dtype, res.shape, res.min(), res.max())
img[res == -1] = (0, 255, 0)
my_show_img(img, 'res', lambda x: cv.cvtColor(x, cv.COLOR_BGR2RGB))

plt.show()
